/*
 * Copyright 2009 XueSong Guo.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package cn.webwheel.compiler;

import javax.tools.FileObject;
import javax.tools.ForwardingJavaFileManager;
import javax.tools.JavaFileObject;
import javax.tools.StandardJavaFileManager;
import java.io.IOException;
import java.net.URISyntaxException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.Map;
import java.util.Set;

public class MemFileManager extends ForwardingJavaFileManager<StandardJavaFileManager> {

    private ClassLoader classLoader;

    private Map<String, MemJavaFileObject> mjfoMap = new HashMap<String, MemJavaFileObject>();

    MemFileManager(StandardJavaFileManager fileManager, ClassLoader classLoader) throws URISyntaxException {
        super(fileManager);
        this.classLoader = classLoader;
    }

    public ClassLoader getClassLoader(Location location) {
        return classLoader;
    }

    public JavaFileObject getJavaFileForOutput(Location location, String className, JavaFileObject.Kind kind, FileObject sibling) throws IOException {
        MemJavaFileObject mjfo;
        try {
            mjfo = new MemJavaFileObject(className);
        } catch (URISyntaxException e) {
            throw new IOException(e);
        }
        mjfoMap.put(className, mjfo);
        return mjfo;
    }

    public Iterable<JavaFileObject> list(Location location, String packageName, Set<JavaFileObject.Kind> kinds, boolean recurse) throws IOException {
        Iterable<JavaFileObject> result = super.list(location, packageName, kinds, recurse);
        if (packageName.equals("cn.webwheel.gen")
                && kinds.contains(JavaFileObject.Kind.CLASS)) {
            ArrayList<JavaFileObject> temp = new ArrayList<JavaFileObject>();
            for (JavaFileObject jfo : result) temp.add(jfo);
            for (Map.Entry<String, MemJavaFileObject> entry : mjfoMap.entrySet()) {
                temp.add(entry.getValue());
            }
            fillMyClassLoader(temp, classLoader);
            result = temp;
        }
        return result;
    }

    private void fillMyClassLoader(ArrayList<JavaFileObject> list, ClassLoader classLoader) {
        if(classLoader==null) return;
        if (classLoader instanceof MyClassLoader) {
            MyClassLoader my = (MyClassLoader) classLoader;
            Set<String> set = my.getLoadedClasses();
            for (String s : set) {
                byte[] code = my.getByteCode(s);
                MemJavaFileObject mjfo;
                try {
                    mjfo = new MemJavaFileObject(s);
                    mjfo.openOutputStream().write(code);
                } catch (Exception e) {
                    throw new RuntimeException("impossible");
                }
                list.add(mjfo);
            }
        }
        fillMyClassLoader(list, classLoader.getParent());
    }

    public String inferBinaryName(Location location, JavaFileObject file) {
        if (file instanceof MemJavaFileObject) {
            return file.getName();
        }
        return super.inferBinaryName(location, file);
    }

    public MemJavaFileObject findJavaFileObject(String className) {
        return mjfoMap.get(className);
    }
}
